# -*- coding: utf-8 -*-
# file: aen.py
# author: songyouwei <youwei0314@gmail.com>
# Copyright (C) 2018. All Rights Reserved.

from layers.dynamic_rnn import DynamicLSTM
from layers.squeeze_embedding import SqueezeEmbedding
from layers.attention import Attention, NoQueryAttention
from layers.point_wise_feed_forward import PositionwiseFeedForward
import torch
import torch.nn as nn
import torch.nn.functional as F

# CrossEntropyLoss for Label Smoothing Regularization,标签正则化
class CrossEntropyLoss_LSR(nn.Module):
    def __init__(self, device, para_LSR=0.2):
        super(CrossEntropyLoss_LSR, self).__init__()
        self.para_LSR = para_LSR
        self.device = device
        self.logSoftmax = nn.LogSoftmax(dim=-1)

    def _toOneHot_smooth(self, label, batchsize, classes):
        prob = self.para_LSR * 1.0 / classes
        one_hot_label = torch.zeros(batchsize, classes) + prob
        for i in range(batchsize):
            index = label[i]
            one_hot_label[i, index] += (1.0 - self.para_LSR)
        return one_hot_label

    def forward(self, pre, label, size_average=True):
        b, c = pre.size()
        one_hot_label = self._toOneHot_smooth(label, b, c).to(self.device)
        loss = torch.sum(-one_hot_label * self.logSoftmax(pre), dim=1)
        if size_average:
            return torch.mean(loss)
        else:
            return torch.sum(loss)

'''
We hereby focus on the bert version.
class AEN_GloVe(nn.Module):
    def __init__(self, embedding_matrix, opt):
        super(AEN_GloVe, self).__init__()
        self.opt = opt
        self.embed = nn.Embedding.from_pretrained(torch.tensor(embedding_matrix, dtype=torch.float))
        self.squeeze_embedding = SqueezeEmbedding()

        self.attn_k = Attention(opt.embed_dim, out_dim=opt.hidden_dim, n_head=8, score_function='mlp', dropout=opt.dropout)
        self.attn_q = Attention(opt.embed_dim, out_dim=opt.hidden_dim, n_head=8, score_function='mlp', dropout=opt.dropout)
        self.ffn_c = PositionwiseFeedForward(opt.hidden_dim, dropout=opt.dropout)
        self.ffn_t = PositionwiseFeedForward(opt.hidden_dim, dropout=opt.dropout)

        self.attn_s1 = Attention(opt.hidden_dim, n_head=8, score_function='mlp', dropout=opt.dropout)

        self.dense = nn.Linear(opt.hidden_dim*3, opt.polarities_dim)

    def forward(self, inputs):
        text_raw_indices, target_indices = inputs[0], inputs[1]
        context_len = torch.sum(text_raw_indices != 0, dim=-1)
        target_len = torch.sum(target_indices != 0, dim=-1)
        context = self.embed(text_raw_indices)
        context = self.squeeze_embedding(context, context_len)
        target = self.embed(target_indices)
        target = self.squeeze_embedding(target, target_len)

        hc, _ = self.attn_k(context, context)
        hc = self.ffn_c(hc)
        ht, _ = self.attn_q(context, target)
        ht = self.ffn_t(ht)

        s1, _ = self.attn_s1(hc, ht)

        context_len = torch.tensor(context_len, dtype=torch.float).to(self.opt.device)
        target_len = torch.tensor(target_len, dtype=torch.float).to(self.opt.device)

        hc_mean = torch.div(torch.sum(hc, dim=1), context_len.view(context_len.size(0), 1))
        ht_mean = torch.div(torch.sum(ht, dim=1), target_len.view(target_len.size(0), 1))
        s1_mean = torch.div(torch.sum(s1, dim=1), context_len.view(context_len.size(0), 1))

        x = torch.cat((hc_mean, s1_mean, ht_mean), dim=-1)
        out = self.dense(x)
        return out
'''


class AEN_BERT(nn.Module):
    def __init__(self, bert, opt):
        super(AEN_BERT, self).__init__()
        self.opt = opt
        self.bert = bert
        self.squeeze_embedding = SqueezeEmbedding()
        self.dropout = nn.Dropout(opt.dropout)

        self.attn_k = Attention(opt.bert_dim, out_dim=opt.hidden_dim, n_head=8, score_function='mlp', dropout=opt.dropout)
        self.attn_q = Attention(opt.bert_dim, out_dim=opt.hidden_dim, n_head=8, score_function='mlp', dropout=opt.dropout)

        self.ffn_c = PositionwiseFeedForward(opt.hidden_dim, dropout=opt.dropout)  #PCT层
        self.ffn_t = PositionwiseFeedForward(opt.hidden_dim, dropout=opt.dropout)

        self.attn_s1 = Attention(opt.hidden_dim, n_head=8, score_function='mlp', dropout=opt.dropout)

        self.dense = nn.Linear(opt.hidden_dim*3, opt.polarities_dim)  #预测情绪极性的分类层

        self.softmax = nn.Softmax(dim=1)

    def forward(self, inputs):
        '''
            在debug的时候，如果加了注释改变了行，但是执行的时候的高亮还是原来的行
           拿到的是： 'text_bert_indices':cls+left+aspect+right+sep
            'aspect_bert_indices':cls+方面+seq
        '''
        context, target = inputs[0], inputs[1]  #拿到输入的tensor,intputs[0],[1]都是16*85，也就是16个句子，每个句子85个单词
        context_len = torch.sum(context != 0, dim=-1)#计算所有句子的长度，以一维张量的形式返回，维度为batchsize
        target_len = torch.sum(target != 0, dim=-1)  #拿到文本长度和方面词长度,方面词是有长度的。返回16维度的向量，cls+单词1+sep=3，cls+单词1+单词2+sep=5

        context = self.squeeze_embedding(context, context_len)  #虽然句子全部补全为85，但是batch里面最长的是60的话，就压缩长度到60
        context, _ = self.bert(context)  #拿到的第一个输出就是词嵌入向量，转换为词嵌入向量计算也花了几秒钟
        context = self.dropout(context)

        target = self.squeeze_embedding(target, target_len) #一开始限制的方面词的长度最大值为多少?
        target, _ = self.bert(target)  #  target:batchsize*TargetSeqlength*dim，第二个参数是？
        target = self.dropout(target)

        hc, _ = self.attn_k(context, context)  #intraMHA，上下文和上下文自己做注意力，输入为(key,query),hc是每个token融合了其他token组成的隐藏变量组成的。输出为(batchsize*SeqLength*hidden dim)
        hc = self.ffn_c(hc)  #输入到PCT层
        ht, _ = self.attn_q(context, target)  #interMHA,context为key，target为query，target对context进行查询
        ht = self.ffn_t(ht)    #注意力编码层？

        s1, _ = self.attn_s1(hc, ht)  #将编码层出来的某个方面词向量和整个上下文向量进行attention
        #dim=1加和，平均池化，每个句子生成一个300维度的代表向量
        hc_mean = torch.div(torch.sum(hc, dim=1), context_len.unsqueeze(1).float())
        ht_mean = torch.div(torch.sum(ht, dim=1), target_len.unsqueeze(1).float())
        s1_mean = torch.div(torch.sum(s1, dim=1), context_len.unsqueeze(1).float())   #平均池化

        x = torch.cat((hc_mean, s1_mean, ht_mean), dim=-1)   #将三个向量连接到一起
        out = self.dense(x)  #CES损失函数自带softmax
        return out
